import os.path as osp
from PIL import Image
from torch.utils import data
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.utils import make_grid
import torchvision.transforms.functional as F
import os
import numpy as np
import torch
import json
import matplotlib.pyplot as plt

from Models.dataloader.miniimagenet.fcn.mini_imagenet import MiniImageNet

data_dir = "datasets"
dataset = "miniimagenet"
image_size = (600,800)

class CustomDataset(Dataset):
    def __init__(self, dataset) -> None:
        super().__init__()
        IMAGE_PATH = os.path.join(data_dir, '{}/images'.format(dataset))
        SPLIT_PATH = os.path.join(data_dir, '{}/split'.format(dataset))
        if dataset in ["miniimagenet", "cub"]:
            setname = "train"
            csv_path = osp.join(SPLIT_PATH, setname + '.csv')
            lines = [x.strip() for x in open(csv_path, 'r').readlines()][1:]
            data = []
            label = []
            lb = -1
            wnids = []
            for l in lines:
                name, wnid = l.split(',')[:2]
                path = osp.join(IMAGE_PATH, name)
                if wnid not in wnids:
                    wnids.append(wnid)
                    lb += 1
                data.append(path)
                label.append(lb)
            self.data = data  # data path of all data
            self.label = label  # label of all data
            self.num_class = len(set(label))

        elif dataset in ["recognition36", "VHR-10"]:
            SPLIT_PATH = os.path.join(data_dir, '{}'.format(dataset))
            if dataset == "recognition36":
                setname = "novel_all"
            else:
                setname = "novel"

            if setname in ["test", "novel_all"]:
                json_path = osp.join(SPLIT_PATH, 'novel_all.json') # 当测试数据集与源数据集不一致时使用
            else:
                json_path = osp.join(SPLIT_PATH, setname + '.json')

            with open(json_path, "r") as f:
                self.meta = json.load(f) # json file(dict) ： {"label_names:[...], "image_names":[...], "image_labels":[...]}

            data = self.meta["image_names"]
            label = self.meta["image_labels"] # dataset 返回的label并无实际作用,只是当做一个类别标识
            data = [_.replace("filelists", "datasets") for _ in data]
    
            self.data = data  # data path of all data
            self.label = label  # label of all data
            self.num_class = len(set(label))

        self.transform = transforms.Compose([
            transforms.Resize(image_size),
            transforms.ToTensor(),])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        path, label = self.data[i], self.label[i]
        image = self.transform(Image.open(path).convert('RGB'))
        return image, label

def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

if __name__ == '__main__':
    plt.rcParams["savefig.bbox"] = 'tight'
    miniimagenet = CustomDataset("VHR-10")
    data_loader = DataLoader(miniimagenet, 12, shuffle=True, num_workers=0)

    image_batches, _ = next(iter(data_loader))
    grid = make_grid(image_batches, 4, 2)
    show(grid)
    plt.show()
